# -*- coding: utf-8 -*-

import os
import json
from datetime import datetime
import numpy as np
import oyaml as yaml

from od.utilities import read_pkg_file


def read_text_label_file(label_file_path):
    with open(label_file_path, "rb") as label_file:
        data = np.frombuffer(label_file.read(), dtype=np.int32)
    x1, x2, y1, y2 = data.tolist()
    w = x2 - x1 + 1
    h = y2 - y1 + 1
    return x1, y1, w, h


def read_json_label_file(label_file_path):
    with open(label_file_path, "r") as label_file:
        data = json.load(label_file)
    image_width = data["image"]["information"]["width"]
    image_height = data["image"]["information"]["height"]
    x1, y1, x2, y2 = [float(e) for e in data["annotation"]["persons"][0]["box"]]
    x1 = int(x1 * image_width)
    y1 = int(y1 * image_height)
    x2 = int(x2 * image_width)
    y2 = int(y2 * image_height)
    w = x2 - x1 + 1
    h = y2 - y1 + 1
    return x1, y1, w, h


def read_label_file(image_path):
    label_file_path = image_path.replace(".jpg", ".txt")
    if os.path.exists(label_file_path):
        return read_text_label_file(label_file_path)
    label_file_path = image_path.replace(".jpg", ".json")
    if os.path.exists(label_file_path):
        return read_json_label_file(label_file_path)
    raise ValueError


def read_image_paths_and_rectangles(folder):
    """
    person_tracker_image 文件夹中存储的图片和对应的Json文件
    Json文件中记录了目标位置信息
    """
    image_names = [name for name in os.listdir(folder) if ".jpg" in name]
    image_paths = [os.path.join(folder, name) for name in image_names]
    rectangles = [read_label_file(image_path) for image_path in image_paths]
    return image_paths, rectangles


def get_to_be_labeled_persons_information(bus_date_folder, total_person_ids, labeled_person_ids):
    to_be_labeled_persons = list()

    to_be_labeled_persons_id = [person_id for person_id in total_person_ids
                                if person_id not in labeled_person_ids]
    reid_pkg_folder = os.path.join(bus_date_folder, "information", "reid_pkgs")
    person_tracker_folder = os.path.join(bus_date_folder, "information", "person_tracker_image")
    for person_id in to_be_labeled_persons_id:
        person_folder = os.path.join(person_tracker_folder, person_id)
        image_paths, rectangles = read_image_paths_and_rectangles(person_folder)
        reid_pkg_file_path = os.path.join(reid_pkg_folder, f"pkg_{int(person_id)}.txt")
        timestamp = read_pkg_file(reid_pkg_file_path)["uodr_end_time"]
        time = datetime.fromtimestamp(int(timestamp)).strftime("%Y-%m-%d %H:%M:%S")
        to_be_labeled_persons.append(
            {
                "person_id": person_id,
                "image": image_paths,
                "rectangle": rectangles,
                "time": time,
                "location": "",
                "index": 0,
            }
        )
    return to_be_labeled_persons


def get_labeled_person_ids(yaml_file_path):
    """
    od_label.yaml 中记录了已经认为匹配的 gallery 和 query 对。
    在 gallery 和 query 对中，记录了对应的图像地址文件夹
    """
    labeled_person_ids = list()
    with open(yaml_file_path, "r", encoding="utf-8") as yaml_file:
        data = yaml.load(yaml_file)
    pairs = data["pairs"]
    for _, pair in pairs.items():
        gallery_person_id = os.path.basename(pair["gallery"]["image"])
        query_person_id = os.path.basename(pair["query"]["image"])
        labeled_person_ids.append(gallery_person_id)
        labeled_person_ids.append(query_person_id)
    return labeled_person_ids


class LabelBackend(object):
    def __init__(self, bus_date_folder):
        self._bus_date_folder = bus_date_folder
        self._label_file_path = os.path.join(
            bus_date_folder, "information", "label", "od_label.yaml")
        self._gallery_persons = self.get_gallery_persons()
        self._query_persons = self.get_query_persons()
        self._gallery_persons_index = 0
        self._query_persons_index = 0

    def get_gallery_persons(self):
        labeled_person_ids = get_labeled_person_ids(self._label_file_path)
        gallery_folder = os.path.join(self._bus_date_folder, "information", "0")
        gallery_person_ids = [name.split(".")[0] for name in os.listdir(gallery_folder) if ".jpg" in name]
        gallery_persons = get_to_be_labeled_persons_information(
            self._bus_date_folder, gallery_person_ids, labeled_person_ids)
        return gallery_persons

    def get_query_persons(self):
        labeled_person_ids = get_labeled_person_ids(self._label_file_path)
        query_folder = os.path.join(self._bus_date_folder, "information", "1")
        query_person_ids = [name.split(".")[0] for name in os.listdir(query_folder) if ".jpg" in name]
        query_persons = get_to_be_labeled_persons_information(
            self._bus_date_folder, query_person_ids, labeled_person_ids)
        return query_persons

    def generate_pair(self, gallery_pop, query_pop):
        gallery = {"image": F"information/person_tracker_image/{gallery_pop['person_id']}",
                   "reid_feature": F"information/reid_pkgs/pkg_{int(gallery_pop['person_id'])}.txt",
                   "time": gallery_pop["time"],
                   "location": ""}
        query = {"image": F"information/person_tracker_image/{query_pop['person_id']}",
                 "reid_feature": F"information/reid_pkgs/pkg_{int(query_pop['person_id'])}.txt",
                 "time": query_pop["time"],
                 "location": ""}
        gallery_feature = read_pkg_file(
            os.path.join(self._bus_date_folder, gallery["reid_feature"]))["uodr_pkg"]
        query_feature = read_pkg_file(
            os.path.join(self._bus_date_folder, query["reid_feature"]))["uodr_pkg"]
        distance = 1 - np.dot(np.array(gallery_feature), np.array(query_feature))
        distance = F"{distance.item():.7f}"
        added_pair = {"gallery": gallery,
                      "query": query,
                      "distance": distance,
                      "source": "manual",
                      "status": "unchecked"}
        return added_pair

    def update_label_file(self, added_pair):
        with open(self._label_file_path, "r", encoding="utf-8") as label_file:
            data = yaml.load(label_file)
        last_key = list(data["pairs"].keys())[-1]
        index = int(last_key.split("_")[-1]) + 1
        key = F"pair_{index:05}"
        data["pairs"][key] = added_pair
        with open(self._label_file_path, "w", encoding="utf-8") as label_file:
            yaml.dump(data, label_file, allow_unicode=True, default_flow_style=False)

    def save_labeled_pair(self):
        """
        标注者点击了"保存并继续"的按钮，保存当前查找到的匹配对
        """
        # remove current gallery and query person
        gallery_persons_index = self.gallery_persons_index
        gallery_pop = self.gallery_persons.pop(gallery_persons_index)
        self.gallery_persons_index = gallery_persons_index % len(self.gallery_persons)
        query_persons_index = self.query_persons_index
        query_pop = self.query_persons.pop(query_persons_index)
        self.query_persons_index = query_persons_index % len(self.query_persons)

        added_pair = self.generate_pair(gallery_pop, query_pop)
        self.update_label_file(added_pair)

    @property
    def gallery_persons(self):
        return self._gallery_persons

    @property
    def gallery_persons_index(self):
        return self._gallery_persons_index

    @gallery_persons_index.setter
    def gallery_persons_index(self, index):
        if 0 <= index < len(self._gallery_persons):
            self._gallery_persons_index = index

    @property
    def gallery_person(self):
        return self._gallery_persons[self._gallery_persons_index]

    @property
    def gallery_person_image_index(self):
        return self.gallery_person["index"]

    @gallery_person_image_index.setter
    def gallery_person_image_index(self, index):
        if 0 <= index < len(self.gallery_person["image"]):
            self.gallery_person["index"] = index

    @property
    def query_persons(self):
        return self._query_persons

    @property
    def query_persons_index(self):
        return self._query_persons_index

    @query_persons_index.setter
    def query_persons_index(self, index):
        if 0 <= index < len(self._query_persons):
            self._query_persons_index = index

    @property
    def query_person(self):
        return self._query_persons[self._query_persons_index]

    @property
    def query_person_image_index(self):
        return self.query_person["index"]

    @query_person_image_index.setter
    def query_person_image_index(self, index):
        if 0 <= index < len(self.query_person["image"]):
            self.query_person["index"] = index
